Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add torchcompile_xentropy executor #1655

Merged
merged 22 commits into from
Feb 9, 2025
Merged

Add torchcompile_xentropy executor #1655

merged 22 commits into from
Feb 9, 2025
+81 −3

Conversation

riccardofelluga
Copy link
Collaborator

What does this PR do?

In an attempt to fix #1654 and partially #1552, this PR adds the necessary ops in the torchcompile_cat list to capture HF CausalLMLoss.

Before merging this PR needs testing with other models, I will write a message with the results of the benchmarks.

@riccardofelluga riccardofelluga mentioned this pull request Jan 17, 2025
Open
@IvanYashchuk
Copy link
Collaborator

This executor is meant for fusions with a cat operation. Cross entropy is usually quite far from RoPE (which uses cat). How does adding cross_entopy to torchcompile_cat help?

@riccardofelluga
Copy link
Collaborator Author

@IvanYashchuk Adding cross entropy does indeed not help with rope but it allows thunder to use a very efficient fused cross entropy triton kernel which perf currently cannot be matched with the other executors(not even apex). From lines 211 and 212, it does seem that while the set of ops was originally though for rope, it is suggested that other ops could have been added in the future making this a fusion executor that comes in rescue of nvFuser when it can improve performance. Do you think that it would be better to create another executor entry for inductor cross entropy only?

@nvMelissa nvMelissa mentioned this pull request Jan 27, 2025
Closed
@riccardofelluga riccardofelluga changed the title [WIP] Add cross_entropy to torchcompile_cat executor [WIP] Add torchcompile_xentropy executor Jan 28, 2025
@riccardofelluga riccardofelluga changed the title [WIP] Add torchcompile_xentropy executor Add torchcompile_xentropy executor Jan 28, 2025
pre-commit-ci bot and others added 11 commits January 28, 2025 15:29
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
284e7d9 
for more information, see https://pre-commit.ci
@riccardofelluga
remove torch.nn.functional.cross_entropy from supported as it is in r…
c4d98e9 
…equired
@ali-alshaar7 @riccardofelluga
bump version (#1708)
9338118
@kshitij12345 @riccardofelluga
nvfuser: add option to allow shape only region (#1702)
c875621
@jjsjann123 @riccardofelluga
Backward transform dependency fix (#1693)
7e857f5
@t-vi @riccardofelluga
pin check schema (#1709)
fcec023
@riccardofelluga
bump transformers version (#1698)
9a40bb0
@riccardofelluga
Reduce test_thunderfx_mistral_nemo_small model size (#1701)
1df5929 
Reducing the models size to speedup CI and make it work in environments with constrained memory size.

I think this change should be fine since with this test we are verifying functionality more than memory footprint. Let me know what you think. 

With this change the peak memory of the test is ~2.6GB instead of ~14GB
@riccardofelluga
fix naming
71bd489
@riccardofelluga
add pad prims
d94a420
@riccardofelluga
add test for the new executor
35c4606
@riccardofelluga riccardofelluga marked this pull request as ready for review January 30, 2025 07:37
@t-vi
Copy link
Collaborator

t-vi commented Jan 30, 2025

I have a hard time seeing another torch.compile executor as the solution here.
How are they supposed to interact?
Will we see consecutive torch compile regions as a result?

@riccardofelluga
Copy link
Collaborator Author

@t-vi
How are they supposed to interact?
Will we see consecutive torch compile regions as a result?

I'd piggyback on this explaination on the differences between torchcompile and torchcompile_cat executors:

lightning-thunder/thunder/executors/torch_compile.py

Lines 209 to 215 in 880419d

# NOTE: [torch_compile_cat_ex vs torch_compile_ex]
# The former only relies on `torch.compile` for the operators where it shines the most and is meant to be used
# together with the nvfuser executor. Its current goal is only to fuse RoPE but the set of ops fused will change as each
# of the fusion backends evolve.
# The latter will try to `torch.compile` all the torch operators and is meant to be used without the nvfuser_executor
# since they would be competing over fusion opportunities. The advantage over simply doing `torch.compile` is that you
# still get all of Thunder's advantages, like enabling custom executors (e.g. with custom triton kernels) before it.

torchcompile_xentropy is treated the same way, it is not supposed to be used with torchcompile. On the other hand i can see the worry since it can be used together with torchcompile_cat, however, the combination of required+supported ops are different from torchcompile_cat so one won't claim the ops for the other.

An issue I see is that we don't strictly enforce this rule, so in case where the user specifies torchcompile and torchcompile_cat or torchcompile_xentropy, the former will claim all the ops

@t-vi
Copy link
Collaborator

t-vi commented Feb 9, 2025 

torchcompile_xentropy is treated the same way, it is not supposed to be used with torchcompile. On the other hand i can see the worry since it can be used together with torchcompile_cat, however, the combination of required+supported ops are different from torchcompile_cat so one won't claim the ops for the other.

My main issue with this is that we should really have one torchcompile fusion executor in order to not split what could be a big fusion. And the solution would be to split the marking of things to be fused by the torchcompile executor and then But let's have that later.

t-vi
t-vi approved these changes Feb 9, 2025
View reviewed changes
Copy link
Collaborator

@t-vi t-vi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @riccardofelluga

@t-vi t-vi merged commit 5b847bc into main Feb 9, 2025
49 checks passed
@t-vi t-vi deleted the torchcompilecat-add-xentropy branch February 9, 2025 10:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@t-vi t-vi t-vi approved these changes

@IvanYashchuk IvanYashchuk Awaiting requested review from IvanYashchuk

@mruberry mruberry Awaiting requested review from mruberry mruberry is a code owner

@lantiga lantiga Awaiting requested review from lantiga lantiga is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

nvFuser using more memory than inductor for HF CausalLMLoss
6 participants
@riccardofelluga @IvanYashchuk @t-vi @ali-alshaar7 @kshitij12345 @jjsjann123